#include <stdio.h>
#include <stdlib.h>
#include <mpi.h>
#include <math.h>
#include "matrix_reader.h"
#include "localize.h"
#include "sparse_blas.h"
#include "util.h"

int main(int argc, char *argv[])
{
    /* local variables */
    int my_rank = 0, p = 0;

    MPI_Comm *common = (MPI_Comm *)malloc(1*sizeof(MPI_Comm));
    common[0] = MPI_COMM_WORLD;

    MPI_Init(&argc, &argv); // initializing MPI environment
    MPI_Comm_rank(MPI_COMM_WORLD, &my_rank); // getting the current rank
    MPI_Comm_size(MPI_COMM_WORLD, &p); // getting number of processes 

    int i = 0, j = 0, k, ss,  k_end = atoi(argv[4]);
       
    /* setting appropriate output environment */
    char stdoutname[128];
    FILE *newstdout;
    sprintf(stdoutname, "./result/out.%d", my_rank);
    newstdout = freopen(stdoutname, "w", stdout);

    if (newstdout == NULL) {
        printf("Failed to redirect stdout to %s.\n", stdoutname);
    }

    ////////////////// CODE BODY /////////////////////////  
    int nnz, my_nnz, nrows;
    int *ia = NULL, *ja = NULL;
    double *A = NULL;
    int nrowsI = 0;
    /* reading the  matrix from a specific file */
    matrix_reader(argv[1], p, my_rank, &nnz, &my_nnz, &nrows, &ia, &ja, &A);
    int modRows = 0;

    /* in the case the rows of the matrix is unable to be divided. */
    if (nrows % p != 0) {
        modRows = p - nrows % p;
    }

    nrowsI = (int)(nrows + modRows) / p;

    /* reading the rhs vector from the input file */
    double *b = (double *)calloc(nrowsI, sizeof(double));
    read_rhs(argv[2], p, my_rank, nrows, nrowsI, b);

    int **receive = NULL;
    int *cntowner= NULL;
    int **to_be_sent= NULL;
    int *cntsent = NULL;

    localize(my_rank, p, nrows, ia, ja, common, &receive, &cntowner, &to_be_sent, &cntsent);

    double t_start = 0, t_end = 0;  
    
    int precond_flag = atoi(argv[5]);
    double gmres_res = 1.e-6; // gmres residual
    int n_total = max_int_array(ja, my_nnz)+1; // ntotal is equal to nlocal+n phantom points
    double norm_r0 = 0.;

    double *u = (double *)calloc(nrowsI, sizeof(double));
    double **v = (double **)malloc(1 * sizeof(double *));
    double **h = (double **)malloc(1 * sizeof(double *));
    double *x0 = (double *)calloc(n_total, sizeof(double));
    double *r0 = (double *)calloc(nrowsI, sizeof(double));
    double *g = (double *)calloc(1, sizeof(double));
    double *c = NULL, *s = NULL, *resi = NULL;
    double delta = 0., gamma = 0.;
    double *X = (double *)calloc(nrowsI , sizeof(double) );

    /* pick the diagonal entries for preconditioning (if it is preconditioned) */
    double *diag = (double *) calloc( nrowsI , sizeof(double) );
    double *yk = (double *) calloc( n_total , sizeof(double) );
    /* check if it be precondioned */
    if(precond_flag) {
        for (i = 0; i < nrowsI; i++) {
            for (j = ia[i]; j < ia[i+1]; j++) {
                if(ja[j] == i) {
                    diag[i] = A[j];
                }
            }
        }
      }
    /* Read the init data from a specific file. */   
    read_rhs(argv[3], p, my_rank , nrows, nrowsI, x0);
    
    /* restarting variables */
    int rst = 0, mm = atoi(argv[6]), completed = 0, total_itr = 1;
       
    t_start = MPI_Wtime();
    /* do restarting */
    for (rst = 0; rst < mm; rst++) {

        /* gather and sync */
        gather_sync(my_rank, p, nrows, ia, ja, receive, cntowner, to_be_sent, cntsent, x0, common);
        /* matrix vector product */
        sparse_matmult(A, x0, r0, ia, ja, nrowsI, common);
        /* updating the final r0 */
        for (i = 0; i < nrowsI; i++) {
            r0[i] = b[i] - r0[i];
        }

        /* normalizing */
        dotproduct(r0, r0, &norm_r0, nrowsI, common);
        norm_r0 = sqrt(norm_r0);
        for (i = 0; i < nrowsI; i++) {
            r0[i] /= norm_r0;
        }
            
        /* initializing the allocation of v */
        h[0] = (double *)calloc( 1 , sizeof(double) );
        v[0] = (double *)calloc( n_total , sizeof(double) );

        for (i = 0; i < nrowsI; i++) {
            v[0][i] = r0[i];
        }

        g[0] = norm_r0;
           
        /* main gmres loop */
        for (k = 0; k < k_end; k++) {
            /* reallocating vars to have enough space for the next items */
            g = (double *) realloc (g,  (k+2) * sizeof(double) );
            g[k+1] = 0.;
            c = (double *) realloc (c,  (k+1) * sizeof(double) );
            s = (double *) realloc (s,  (k+1) * sizeof(double) );
            resi = (double *) realloc (resi,  (k+1) * sizeof(double) );
      
            // reallocating v and h
            v = (double **)realloc (v , (k+2) * sizeof(double *) );
            v[k+1] = (double *)calloc( n_total , sizeof(double) );
            h = (double **)realloc (h , (k+2) * sizeof(double *) );
            h[k+1] = (double *)calloc( 1 , sizeof(double) );
      
            for (j = 0; j <= k + 1; j++) {
                h[j] = (double *)realloc(h[j], (k+1) *sizeof(double));
            }

            /* NOTE : THE FOLLOWING IF STATEMENT INCREASES PERFORMANCE WHEN
             *        IT IS NOT SUPPOSED TO BE PRECONDITIONED BECAUSE IT
             *        JUST SKIPS THE INITIALIZATION OF VECTOR [yk].
             */
            if (precond_flag) {
                for (i = 0; i < nrowsI; i++) {
                    if (diag[i]) {
                        yk[i] = v[k][i] / diag[i];
                    } else {
                        yk[i] = v[k][i];
                    }
                }

                gather_sync(my_rank, p, nrows, ia, ja, receive, cntowner, to_be_sent, cntsent, yk, common);
                sparse_matmult(A, yk, u, ia, ja, nrowsI, common);
            } else {
                gather_sync(my_rank, p, nrows, ia, ja, receive, cntowner, to_be_sent, cntsent, v[k], common);
                sparse_matmult(A, v[k], u, ia, ja, nrowsI, common);
            }

            for (j = 0; j <= k; j++) {
                dotproduct(v[j], u, &h[j][k], nrowsI, common);
                for(ss = 0; ss < nrowsI; ss++) {
                    u[ss] -= h[j][k] * v[j][ss];
                }
            }
            
            dotproduct(u, u, &h[k+1][k], nrowsI, common); // dot product
            h[k+1][k] = sqrt(h[k+1][k]); // norm
            
            /* updating v[k+1] */
            for(ss = 0; ss < nrowsI; ss++) {
                v[k+1][ss] = u[ss] / h[k+1][k];
            }

            for (j = 0; j < k; j++) {
                delta = h[j][k];
                h[j][k] = c[j] * delta + s[j] * h[j+1][k];
                h[j+1][k] = -s[j] * delta + c[j] * h[j+1][k];
            }

            gamma = sqrt(h[k][k] * h[k][k] + h[k+1][k]*h[k+1][k]);
            c[k] = h[k][k] / gamma;
            s[k] = h[k+1][k] / gamma;
            h[k][k] = gamma;
            h[k+1][k] = 0.;
            delta = g[k];
            g[k] = c[k] * delta + s[k] * g[k+1];
            g[k+1] = -s[k] * delta + c[k] * g[k+1];
            resi[k] = fabs(g[k+1]);

            if (resi[k] <= gmres_res) {
                completed = 1;
                break;
            }
        /* end of the main gmres loop */
        }

        /* terminate restart loop ASAP */
        if (completed == 0) {
            k--; // terminate restart loop ASAP!
        }
      
        /* calculate  and allocate alpha */
        double *alpha = (double *)calloc(k+1 , sizeof(double));
        for (j = k ; j >= 0; j--) {
            alpha[j] = g[j]/h[j][j];
            for (ss = (j+1) ; ss <= k; ss++) {
                alpha[j] -= (h[j][ss]/h[j][j] * alpha[ss]);
            }
        }

        /* compute zk */
        double *zk = (double *)calloc(nrowsI , sizeof(double) );
        for (j = 0; j <= k; j++) {
            for (ss = 0; ss < nrowsI; ss++) {
                zk[ss] += alpha[j] * v[j][ss];
            }
        }

        if(precond_flag) {
            for (i = 0; i < nrowsI; i++ ) {
                if (diag[i]) {
                    zk[i] /= diag[i];
                }
            }
        }
           
        /* compute solution */
        for (ss = 0; ss < nrowsI; ss++) {
            X[ss] = x0[ss] + zk[ss];
        }
        /* report the solution
         * preparinf for restart operation ...
         * put x into x0
         */
        for( i = 0; i < nrowsI; i++) {
            x0[i] = X[i];
        }
      
        /* clean ups */
        free(alpha);
        free(zk);
        for (j = 0 ; j <= (k+1); j++) {
            free(v[j]);
            free(h[j]);
        }
      
    } /* end of restarting loop */
    /* put a time mark */
    t_end = MPI_Wtime();
    if (my_rank != p - 1) {
        for (i = 0; i < nrowsI; i++) {
            fprintf(stdout, "%.8e\n", X[i]);
        }
    } else {
        for (i = 0; i < nrowsI - modRows; i++) {
            fprintf(stdout, "%.8e\n", X[i]);
        }
    }
    printf("The program iteration turns is : %u\n", k);
    printf("The program calculating time is : %f seconds.\n", (t_end - t_start));

    

    ///////////////END OF CODE BODY //////////////////////
       
    /* closing outputs, finishing MPI session */
    fclose(newstdout);
    MPI_Finalize();


    /* completed successfully! */
    return 0;

}

